# 数据处理
import pandas as pd
from gensim import utils	
import gensim.parsing.preprocessing as gsp
from sklearn.model_selection import train_test_split
# TF-IDF
from sklearn.feature_extraction.text import TfidfVectorizer


# 数据预处理
def pretreat(p):
	filters = [	
		gsp.strip_tags,					# 去除标签（如<html>, <p>等）
		gsp.strip_punctuation,			# 去除标点
		gsp.strip_multiple_whitespaces,	# 去除多余空格
		gsp.strip_numeric,				# 去除整数、数字
		gsp.remove_stopwords,			# 去除停用词（如and、to、the等）
		gsp.strip_short,				# 英文字母小写化
		gsp.stem_text					# 词干提取（将单词转换至词源形式）
		]
	
	p = p.lower()
	p = utils.to_unicode(p)
	for f in filters:
		p = f(p)
	
	return p

# 标签转数字
def label2num(l):
	label_enum = {
		"tech":0,
		"business":1,
		"sport":2,
		"entertainment":3,
		"politics":4 
		}
	
	return label_enum[l]

# 划分训练集和测试集
def devide(bbc_df):
	X_train, X_test, y_train, y_test = train_test_split(bbc_df["news"], bbc_df["type"], train_size=200, test_size=100)
	
	X_train = X_train.apply(lambda X: pretreat(X))
	X_test = X_test.apply(lambda X: pretreat(X))
	y_train = y_train.apply(lambda y: label2num(y))
	y_test = y_test.apply(lambda y: label2num(y))
	
	return X_train, X_test, y_train, y_test

# TF-IDF特征选择
def TFIDF(X_train, X_test, y_train, y_test):
	bbc_tf = TfidfVectorizer(stop_words="english", min_df=0.05)		# 转化为词频矩阵 english=去除英语停用词 去除词频小于0.05
	# 训练 构建词汇表以及词项idf值
	X_train = bbc_tf.fit_transform(X_train)
	X_train_tfidf = X_train.toarray()
	#print(bbc_tf.get_feature_names())
	
	# 对测试集进行tf-idf权重计算
	X_test = bbc_tf.transform(X_test)
	X_test_tfidf = X_test.toarray()			# 测试集TF-IDF权重矩阵
	
	return X_train_tfidf, X_test_tfidf, y_train, y_test